import random
import matplotlib.pyplot as plt
import numpy as np

class TDMA_Node:
    def __init__(self, id, max_distance, num_slots):
        self.id = id
        self.max_distance = max_distance
        self.num_slots = num_slots
        self.success_probabilities = self.calculate_success_probabilities()

    def calculate_success_probabilities(self):
        # 计算每个时隙内节点的成功发送概率，距离越远成功概率越低
        success_probabilities = []
        for slot in range(self.num_slots):
            success_prob=0.4
            # distance = random.randint(1, self.max_distance)
            # if distance <= 5:
            #     success_prob = 0.5
            # elif distance <= 10:
            #     success_prob = 0.3
            # else:
            #     success_prob = 0.1
            success_probabilities.append(success_prob)
        return success_probabilities

class TDMA_Simulation:
    def __init__(self, num_nodes, num_slots, max_distance):
        self.num_nodes = num_nodes
        self.num_slots = num_slots
        self.max_distance = max_distance
        self.nodes = [TDMA_Node(i, max_distance, num_slots) for i in range(num_nodes)]

    def simulate(self):
        throughputs = []  # 存储吞吐量

        for slot in range(self.num_slots):
            successful_transmissions_count = 0

            # 模拟每个时隙内节点的数据传输
            for node in self.nodes:
                success_prob = node.success_probabilities[slot]
                if random.random() < success_prob:
                    successful_transmissions_count += 1

            # 计算当前时隙的吞吐量
            throughput = successful_transmissions_count / self.num_nodes
            throughputs.append(throughput)

        # 使用移动平均平滑吞吐量数据
        smoothed_throughputs = self.smooth_data(throughputs, window_size=1000)

        return smoothed_throughputs

    def smooth_data(self, data, window_size):
        smoothed_data = []
        for i in range(len(data)):
            start_index = max(0, i - window_size + 1)
            end_index = i + 1
            window_data = data[start_index:end_index]
            smoothed_value = np.mean(window_data)
            smoothed_data.append(smoothed_value)
        return smoothed_data

    def plot_throughput(self, throughputs):
        # 绘制平滑后的吞吐量随时间变化的折线图
        plt.plot(range(1, self.num_slots + 1), throughputs)
        plt.xlabel('Time Slots')
        plt.ylabel('Throughput')
        plt.title('Smoothed Throughput over Time in TDMA Simulation')
        plt.ylim(0, 1)  # 设置纵轴范围为0到1
        #plt.show()

def run_tdma_simulation(num_nodes, num_slots, max_distance):
    tdma_simulation = TDMA_Simulation(num_nodes, num_slots, max_distance)
    throughputs = tdma_simulation.simulate()
    np.savetxt('rewards/only_TDMA.txt', throughputs)
    tdma_simulation.plot_throughput(throughputs)

# 如果需要在其他模块中调用 TDMA 模拟函数，可以使用以下方式：
# from TDMA import run_tdma_simulation

# 调用示例：
run_tdma_simulation(num_nodes=5, num_slots=10000, max_distance=20)
